"""
LangChain Agent 应用
这是一个使用LangChain构建的智能Agent应用，支持工具调用、记忆管理和对话功能
"""

from langchain.agents import initialize_agent, AgentType, Tool
from langchain.memory import ConversationBufferMemory
from langchain_ollama import OllamaLLM
from datetime import datetime


class CalculatorTool:
    """简单的计算器工具"""
    
    def add(self, a: float, b: float) -> str:
        """加法运算"""
        result = a + b
        return f"计算结果: {a} + {b} = {result}"
    
    def subtract(self, a: float, b: float) -> str:
        """减法运算"""
        result = a - b
        return f"计算结果: {a} - {b} = {result}"
    
    def multiply(self, a: float, b: float) -> str:
        """乘法运算"""
        result = a * b
        return f"计算结果: {a} × {b} = {result}"
    
    def divide(self, a: float, b: float) -> str:
        """除法运算"""
        if b == 0:
            return "错误: 除数不能为零"
        result = a / b
        return f"计算结果: {a} ÷ {b} = {result}"


class WeatherTool:
    """天气查询工具（模拟）"""
    
    def __init__(self):
        self.weather_data = {
            "北京": "晴天, 温度: 25°C, 湿度: 60%",
            "上海": "多云, 温度: 22°C, 湿度: 70%",
            "深圳": "雨天, 温度: 28°C, 湿度: 80%",
            "杭州": "晴天, 温度: 24°C, 湿度: 65%",
        }
    
    def get_weather(self, city: str) -> str:
        """获取指定城市的天气信息"""
        city = city.strip()
        if city in self.weather_data:
            return f"{city}的天气: {self.weather_data[city]}"
        else:
            available_cities = ", ".join(self.weather_data.keys())
            return f"抱歉，目前没有{city}的天气数据。可用城市: {available_cities}"


class TimeTool:
    """时间查询工具"""
    
    def get_current_time(self) -> str:
        """获取当前时间"""
        now = datetime.now()
        return f"当前时间: {now.strftime('%Y-%m-%d %H:%M:%S')}"
    
    def get_date(self) -> str:
        """获取当前日期"""
        now = datetime.now()
        return f"今天是: {now.strftime('%Y年%m月%d日 %A')}"


def main():
    """主函数"""
    print("=" * 60)
    print("🤖 LangChain Agent 应用启动")
    print("=" * 60)
    print("\n此Agent支持以下功能:")
    print("1. 数学计算 (加减乘除)")
    print("2. 天气查询")
    print("3. 时间查询")
    print("4. 自然语言对话")
    print("\n输入 'quit' 或 'exit' 退出程序\n")
    
    # 初始化LLM (使用Ollama本地模型)
    try:
        llm = OllamaLLM(model="qwen", temperature=0.7)
        print("✅ LLM模型加载成功\n")
    except Exception as e:
        print(f"❌ LLM模型加载失败: {e}")
        print("请确保已安装并运行Ollama服务")
        print("运行命令: ollama pull qwen")
        return
    
    # 创建工具实例
    calc = CalculatorTool()
    weather = WeatherTool()
    time_tool = TimeTool()
    
    # 定义工具列表
    tools = [
        Tool(
            name="Calculator",
            func=lambda query: "请提供计算表达式",
            description="用于执行基本的数学运算（加法、减法、乘法、除法）"
        ),
        Tool(
            name="WeatherTool",
            func=weather.get_weather,
            description="用于查询指定城市的天气信息，输入城市名称即可"
        ),
        Tool(
            name="TimeTool",
            func=lambda x: time_tool.get_current_time() if "时间" in x else time_tool.get_date(),
            description="用于查询当前时间和日期"
        ),
    ]
    
    # 创建记忆
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        return_messages=True
    )
    
    # 初始化Agent
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
        memory=memory,
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=3,
        early_stopping_method="generate"
    )
    
    # 交互循环
    conversation_count = 0
    while True:
        try:
            # 获取用户输入
            user_input = input("\n👤 您: ").strip()
            
            if not user_input:
                continue
            
            # 检查退出命令
            if user_input.lower() in ['quit', 'exit', '退出', 'q']:
                print("\n👋 再见，期待下次见面！")
                break
            
            # 执行Agent推理
            print("\n🤖 Agent思考中...\n")
            conversation_count += 1
            
            response = agent.run(input=user_input)
            
            print(f"\n🤖 Agent: {response}")
            
        except KeyboardInterrupt:
            print("\n\n⚠️  程序被用户中断")
            break
        except Exception as e:
            print(f"\n❌ 发生错误: {e}")
            print("请重试或输入 'quit' 退出")
    
    # 显示统计信息
    print(f"\n📊 本次会话统计:")
    print(f"   - 对话轮数: {conversation_count}")
    print(f"   - 记忆状态: {'已保存' if memory else '未保存'}")


if __name__ == "__main__":
    main()

